package trie;

import java.util.TreeMap;

/**
 * @program: play-structure
 * @author: baichen
 * 字典树,本质是多叉树,存放的是从根结点都叶子结点组成的单词数目，即单词数目
 * 力扣208
 * 可以理解为一个字符一个节点
 **/
public class Trie {
    private class Node {
        public boolean isWord;      //判断当前的节点是否代表一个单词的结尾
        // 到下一个节点的映射，字符串是Character,可以根据实际情况替换为别的数据类型
        // 从Character到Node，从字符映射到节点
        public TreeMap<Character, Node> next;

        public Node(boolean isWord) {
            this.isWord = isWord;
            next = new TreeMap<>();
        }

        public Node() {
            this(false);    //大多数当前节点都不是单词结尾
        }
    }

    private Node root;
    private int size;

    public Trie() {
        root = new Node();
        size = 0;
    }

    // 获得Trie中存储的单词数量
    public int getSize() {
        return size;
    }

    // 向Trie中添加一个新的单词word
    public void add(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            //这里是使用了映射，字符映射到节点，所以根据字符查找这个节点是否存在
            //如果为空(即不存在)则存入这个字符并新建一个节点
            if (cur.next.get(c) == null) {
                cur.next.put(c, new Node());
            }
            //如果存在则直接指向这个节点，新增这个字符对应的节点
            cur = cur.next.get(c);
        }
        //不添加重复的单词
        if (!cur.isWord) {  //该节点不是任何单词的结尾,即添加的单词不存在
            cur.isWord = true;
            size++;     //添加了一个单词，所以size要加一
        }
    }

    // 查询单词word是否在Trie中
    public boolean contains(String word) {
        Node cur = root;
        for (int i = 0; i < word.length(); i++) {
            char c = word.charAt(i);
            if (cur.next.get(c) == null)    //cur 是否包含 c 到下一个节点的映射
                return false;
            cur = cur.next.get(c);  //cur包含 c 到下一个节点的映射
        }
        //这里不能直接返回true，比如在panda中查找pan，pan并不存在，所以通过判断是否是结尾即可
        return cur.isWord;
    }

    // 查询是否在Trie中有单词以prefix为前缀
    public boolean isPrefix(String prefix) {
        Node cur = root;
        for (int i = 0; i < prefix.length(); i++) {
            char c = prefix.charAt(i);
            if (cur.next.get(c) == null)    //cur 是否包含 c 到下一个节点的映射
                return false;
            cur = cur.next.get(c);  //cur包含 c 到下一个节点的映射
        }
        //这里不能直接返回cur.isWord，比如在panda中查找pan，pan存在，但并不是单词结尾
        // 所以直接返回true即可
        return true;
    }
}
